# -*- coding: utf-8 -*-
"""
GraphRAG Schema - 简化版（基于LangChain）
"""
from typing import Optional, Dict, Any, List
from pydantic import BaseModel, Field


class GraphRAGRequest(BaseModel):
    """GraphRAG查询请求"""
    question: str = Field(
        ..., 
        description="用户自然语言问题",
        example="海豹车型在2022年到2024年的销量趋势是什么？"
    )
    
    return_cypher: Optional[bool] = Field(
        default=True,
        description="是否返回生成的Cypher查询"
    )
    
    return_context: Optional[bool] = Field(
        default=True,
        description="是否返回图谱检索的原始数据"
    )


class GraphRAGResponse(BaseModel):
    """GraphRAG查询响应"""
    answer: str = Field(
        ...,
        description="Markdown格式的分析报告"
    )
    
    cypher: Optional[str] = Field(
        default=None,
        description="LLM生成的Cypher查询"
    )
    
    context: Optional[List[Dict[str, Any]]] = Field(
        default=None,
        description="从图谱检索的原始数据"
    )
    
    metadata: Dict[str, Any] = Field(
        default={},
        description="元数据（问题、模型等）"
    )
    
    error: Optional[str] = Field(
        default=None,
        description="错误信息（如果有）"
    )


class GraphStatsResponse(BaseModel):
    """图谱统计响应"""
    total_nodes: int = Field(..., description="总节点数")
    total_relationships: int = Field(..., description="总关系数")
    node_statistics: Dict[str, int] = Field(default={}, description="各类型节点数量")
    relationship_statistics: Dict[str, int] = Field(default={}, description="各类型关系数量")
    schema: Optional[str] = Field(default=None, description="图谱Schema")


class HealthCheckResponse(BaseModel):
    """健康检查响应"""
    status: str = Field(..., description="连接状态: connected/error")
    test_result: Optional[List[Dict[str, Any]]] = Field(default=None, description="测试查询结果")
    schema_loaded: Optional[bool] = Field(default=None, description="Schema是否加载")
    error: Optional[str] = Field(default=None, description="错误信息")
